#!/usr/bin/env python3
# -*- encoding: utf-8 -*-
'''
@File         :pos_control_offb.py
@Description  :位置环控制，主要用于领机的offb控制
@Time         :2024/05/30 16:54:00
@Author       :Lin Yuheng
@Version      :1.0
'''

import rospy
import tty, termios
import sys, select
from mavros_msgs.msg import AttitudeTarget, PositionTarget, State
from mavros_msgs.srv import CommandBool, SetMode
from geometry_msgs.msg import Pose, PoseStamped, Twist, TwistStamped
from nav_msgs.msg import Odometry
from std_msgs.msg import String
import time
import PyKDL
import math
import numpy as np

from outdoor_offb.msg import *

class Controller:
    def __init__(self):
        self.plane_id = 0
        self.leader_id = "plane_0"
        self.vehicle_id = "plane_1"
        self.px = 0
        self.py = 0
        self.pz = 0
        self.vx, self.vy, self.vz = 0, 0, 0
        self.roll, self.pitch, self.yaw = 0, 0, 0
        self.phi_dot, self.theta_dot, self.psi_dot = 0, 0, 0
        
        self.leader_pos_n = 0
        self.leader_pos_e = 0
        self.leader_pos_d = 0
        
        self.vehicle_state = State()
        self.should_exit = False

        # subscribers
        self.local_pose_sub = rospy.Subscriber(
            self.vehicle_id + "/mavros/local_position/pose",
            PoseStamped,
            callback=self.local_pose_callback,
        )
        self.mavros_sub = rospy.Subscriber(
            self.vehicle_id + "/mavros/state",
            State,
            callback=self.mavros_state_callback,
        )
        
        # 【订阅】虚拟领机的状态量
        self.leader_states_sub = rospy.Subscriber(
            self.leader_id + "/fixed_wing_control/leader_state",
            Leaderstates,
            callback=self.leader_states_callback
        )

        # pub
        self.target_motion_pub = rospy.Publisher(
            self.vehicle_id + "/mavros/setpoint_position/local",
            PoseStamped,
            queue_size=10
        )

        self.target_raw_pose = PoseStamped()

    def local_pose_callback(self, msg: PoseStamped):
        self.px = msg.pose.position.x
        self.py = msg.pose.position.y
        self.pz = msg.pose.position.z

    # def local_vel_callback(self, msg: TwistStamped):
    #     self.vx = msg.twist.linear.x
    #     self.vy = msg.twist.linear.y
    #     self.vz = msg.twist.linear.z

    # def odom_callback(self, msg: Odometry):
    #     x = msg.pose.pose.orientation.x
    #     y = msg.pose.pose.orientation.y
    #     z = msg.pose.pose.orientation.z
    #     w = msg.pose.pose.orientation.w
    #     rot = PyKDL.Rotation.Quaternion(x, y, z, w)
    #     self.roll = rot.GetRPY()[0]
    #     self.pitch = rot.GetRPY()[1]
    #     self.yaw = rot.GetRPY()[2]

    #     self.phi_dot = msg.twist.twist.angular.x
    #     self.theta_dot = msg.twist.twist.angular.y
    #     self.psi_dot = msg.twist.twist.angular.z
        
    def leader_states_callback(self, msg:Leaderstates):
        self.leader_pos_n = msg.ned_pos_x
        self.leader_pos_e = msg.ned_pos_y
        self.leader_pos_d = msg.ned_pos_z
        
    def mavros_state_callback(self, msg: State):
        if self.vehicle_state.mode == State.MODE_PX4_OFFBOARD:
            if msg.mode != State.MODE_PX4_OFFBOARD:
                rospy.logwarn("Vehicle quit offboard mode, quitting!")
                self.should_exit = True
        self.vehicle_state = msg
        
    def distance_to_setpoint(self, setpoint):
        return math.sqrt((self.px - setpoint[0]) ** 2 + (self.py - setpoint[1]) ** 2)

    def start(self):
        print("start fixed-wing path control")
        rospy.init_node(
            self.vehicle_id + "_path_control_node"
        )
        rate = rospy.Rate(20)
        
        waypoint_start = [50, -100]
        waypoint = waypoint_start

        while not rospy.is_shutdown():
            if self.should_exit:
                break
            else:
                # /mavros/setpoint_position/local 为enu坐标系
                waypoint[0] = self.leader_pos_e
                waypoint[1] = self.leader_pos_n
                upward = abs(self.leader_pos_d)
                
                dist = self.distance_to_setpoint(waypoint)

                self.target_raw_pose.pose.position.x = waypoint[0]
                self.target_raw_pose.pose.position.y = waypoint[1]
                self.target_raw_pose.pose.position.z = upward

                try:
                    self.target_motion_pub.publish(self.target_raw_pose)
                except rospy.ROSException as e:
                    pass
                rate.sleep()

                rospy.loginfo(
                    "wp:(%f,%f); position: (%f, %f); distance:%f", waypoint[1], waypoint[0], self.py, self.px, dist
                )


if __name__ == "__main__":
    Con = Controller()
    try:
        Con.start()
    except rospy.ROSInterruptException:
        pass
